Skip to content

Conversation

@DeadLining
Copy link

Fix Model Parallelism Error: KVCache CUDA Device Switching Issue

The RuntimeError:
Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument tensors in method wrapper_CUDA_cat)
error mentioned in #issues8 was caused by KVCache failing to switch CUDA devices correctly during inference.

This change is intended to fix the issue by ensuring proper device switching during inference.

@krishnateja95
Copy link

Thanks for the PR. I am able to run across multiple GPUs now.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants